package com.example.maze.util;

import java.util.HashMap;
import java.util.Map;

/**
 * 并查集（Union-Find）数据结构实现
 * 用于Kruskal算法生成迷宫
 * 
 * @author maze-demo
 */
public class UnionFind<T> {
    private final Map<T, T> parent;
    private final Map<T, Integer> rank;

    public UnionFind() {
        this.parent = new HashMap<>();
        this.rank = new HashMap<>();
    }

    /**
     * 创建一个新的集合，包含元素x
     */
    public void makeSet(T x) {
        if (!parent.containsKey(x)) {
            parent.put(x, x);
            rank.put(x, 0);
        }
    }

    /**
     * 查找元素x所属集合的代表元素（根节点）
     * 使用路径压缩优化
     */
    public T find(T x) {
        if (!parent.containsKey(x)) {
            makeSet(x);
        }
        
        T root = parent.get(x);
        if (!root.equals(x)) {
            // 路径压缩：将节点直接连接到根节点
            parent.put(x, find(root));
        }
        return parent.get(x);
    }

    /**
     * 合并包含元素x和y的两个集合
     * 使用按秩合并优化
     */
    public boolean union(T x, T y) {
        T rootX = find(x);
        T rootY = find(y);

        // 如果已经在同一个集合中，返回false
        if (rootX.equals(rootY)) {
            return false;
        }

        // 按秩合并：将较小的树连接到较大的树下
        int rankX = rank.get(rootX);
        int rankY = rank.get(rootY);

        if (rankX < rankY) {
            parent.put(rootX, rootY);
        } else if (rankX > rankY) {
            parent.put(rootY, rootX);
        } else {
            parent.put(rootY, rootX);
            rank.put(rootX, rankX + 1);
        }

        return true;
    }

    /**
     * 检查两个元素是否在同一个集合中
     */
    public boolean connected(T x, T y) {
        return find(x).equals(find(y));
    }

    /**
     * 获取集合数量
     */
    public int getSetCount() {
        Map<T, Integer> rootCount = new HashMap<>();
        for (T element : parent.keySet()) {
            T root = find(element);
            rootCount.put(root, rootCount.getOrDefault(root, 0) + 1);
        }
        return rootCount.size();
    }
}